
import os
import sys
sys.path.append('..')
import logging
logging.basicConfig(level=logging.ERROR)
from paddleocr import PaddleOCR,draw_ocr
from PIL import Image, ImageDraw
import config
import time
from utils import dict_json_saver as udjs
# Paddleocr supports Chinese, English, French, German, Korean and Japanese.
# You can set the parameter `lang` as `ch`, `en`, `fr`, `german`, `korean`, `japan`
# to switch the language model in order.

# ocr = PaddleOCR(use_gpu=False, use_angle_cls=True, lang='ch',enable_mkldnn=True) # need to run only once to download and load model into memory

ocr = PaddleOCR(
    use_gpu=False, 
    use_angle_cls=False, 
    lang='ch', 
    enable_mkldnn=True,
    rec_model_dir=os.path.join(config.ROOT_PATH, 'utils/ch_PP-OCRv4_rec'), # ch_PP-ch_PP-OCRv3_rec_infer enable_mkldnn=True 速度不错！
    )



def ocr_img(src_p=os.path.join(config.ROOT_PATH, 'temp/screenshot.png')):
    t0 = time.time()
    
    with Image.open(src_p) as img:
        # 确保图片格式为RGB
        if img.mode in ["RGBA", "LA"] or (img.mode == "P" and "transparency" in img.info):
            # 使用白色背景创建一个新图像
            background = Image.new("RGB", img.size, (255, 255, 255))
            background.paste(img, mask=img.split()[3])  # 3是alpha通道
            img = background

        draw = ImageDraw.Draw(img)
    result = ocr.ocr(src_p, cls=True)
    result = result[0]
    boxes = [line[0] for line in result]
    texts = [line[1][0] for line in result]  # 所有识别的文字
    find_list = []

    scores = [line[1][1] for line in result]
    tlist =[]
    res_dict = {}

    for ind,text in enumerate(texts):
        box = boxes[ind]

        # 计算矩形的左上角和右下角坐标
        x_coordinates = [point[0] for point in box]
        y_coordinates = [point[1] for point in box]
        left = min(x_coordinates)
        top = min(y_coordinates)
        right = max(x_coordinates)
        bottom = max(y_coordinates)

        center_x = sum(x_coordinates) / len(x_coordinates)
        center_y = sum(y_coordinates) / len(y_coordinates)
        res_dict[text] = (center_x, center_y)
        # 绘制矩形框
        draw.rectangle([left, top, right, bottom], outline="red", width=2)

    # if __name__=='__main__':
    #     img.save(os.path.join(config.ROOT_PATH, 'temp/screenshot_rec.jpg'), format='JPEG')
    #     udjs.save_dict_to_json(res_dict, file_path=os.path.join(config.ROOT_PATH, 'temp/screenshot_rec.json'))
    #     print(round(time.time()-t0, 3))
    img.save(os.path.join(config.ROOT_PATH, 'temp/screenshot_rec.jpg'), format='JPEG')
    udjs.save_dict_to_json(res_dict, file_path=os.path.join(config.ROOT_PATH, 'temp/screenshot_rec.json'))
    print(round(time.time()-t0, 3))

        
    return res_dict


def run():
    ocr_img()

if __name__=='__main__':
    run()